#include "aescrypto.h"

AESCrypto::AESCrypto(unsigned char* key) {
    memcpy(this->key, key, AES_BLOCK_SIZE);
}
std::string AESCrypto::encrypt(const std::string& plaintext) {
    // 生成随机IV
    unsigned char iv[AES_BLOCK_SIZE];
    RAND_bytes(iv, AES_BLOCK_SIZE);

    // 填充明文
    int paddedLen = plaintext.length() + (AES_BLOCK_SIZE - (plaintext.length() % AES_BLOCK_SIZE));
    unsigned char* paddedData = new unsigned char[paddedLen];
    memcpy(paddedData, plaintext.c_str(), plaintext.length());
    PKCS7_padding(paddedData, plaintext.length(), AES_BLOCK_SIZE);

    // 加密
    unsigned char* ciphertext = new unsigned char[paddedLen + AES_BLOCK_SIZE];
    memcpy(ciphertext, iv, AES_BLOCK_SIZE); // 将IV放在密文前面

    AES_KEY aesKey;
    AES_set_encrypt_key(key, 128, &aesKey);
    AES_cbc_encrypt(paddedData, ciphertext + AES_BLOCK_SIZE, paddedLen, &aesKey, iv, AES_ENCRYPT);

    std::string result((char*)ciphertext, paddedLen + AES_BLOCK_SIZE);
    delete[] paddedData;
    delete[] ciphertext;
    return result;
}
std::string AESCrypto::decrypt(const std::string& ciphertext) {
    if (ciphertext.length() <= AES_BLOCK_SIZE) {
        throw std::runtime_error("Invalid ciphertext length");
    }

    // 提取IV
    unsigned char iv[AES_BLOCK_SIZE];
    memcpy(iv, ciphertext.c_str(), AES_BLOCK_SIZE);

    // 解密
    int cipherLen = ciphertext.length() - AES_BLOCK_SIZE;
    unsigned char* plaintext = new unsigned char[cipherLen];

    AES_KEY aesKey;
    AES_set_decrypt_key(key, 128, &aesKey);
    AES_cbc_encrypt((unsigned char*)ciphertext.c_str() + AES_BLOCK_SIZE, plaintext, cipherLen, &aesKey, iv, AES_DECRYPT);

    // 去除填充
    int actualLen = PKCS7_unpadding(plaintext, cipherLen);

    std::string result((char*)plaintext, actualLen);
    delete[] plaintext;
    return result;
}
void AESCrypto::PKCS7_padding(unsigned char* data, int len, int blockSize){
    int pad = blockSize - (len % blockSize);
    for (int i = len; i < len + pad; i++) {
        data[i] = pad;
    }

}
int AESCrypto::PKCS7_unpadding(unsigned char* data, int len) {
    int pad = data[len - 1];
    if (pad > AES_BLOCK_SIZE) return len; // 无效填充
    return len - pad;

}
